from __future__ import division # treat all division like floating point
"""
Must supply a comma separated input-file with lines consisting of
    Src, Dest, Volume

    Ex:
    0, 1, 100
    1, 2, 120
    2, 1, 50
    2, 0, 80

In the above example, exactly 100 passengers will show up wanting to travel from
station 0 to station 1. No passengers will want to travel from station 1 to
staition 0, since no line specifies that connection. You may also use a tab-separated
input file.
"""
import sys
from optparse import OptionParser
import random
from collections import defaultdict
import warnings

optpar = OptionParser(usage="usage: %prog [options] input-file")
optpar.add_option("-t", "--time-distribution", action="store", nargs=2, dest="beta",
                  default=(1,1), type='float',
                  help = "Two shape parameters for a beta distribution "\
                  "controlling the passenger start times. Default is 1 1 "\
                  "(uniform)."
                  )
optpar.add_option("-e", "--endtime", default=60.0, type='float',
                  help="Latest start time allowed for any passenger.")
optpar.add_option("-o", "--output-file", default='-', dest="out",
                  help="Filename for output. Use '-' for stdout (default).")

options, args = optpar.parse_args()

if len(args) != 1:
    print optpar.error("Must supply an input-file name.")
inputfile = open(args[0], 'Ur')
if options.out == '-':
    outfile = sys.stdout
else:
    outfile = open(options.out, 'w')

# parse the input
trips = defaultdict(int) # keyed by (origin, dest) tuples, values are trip counts.
stations = set() # set of unique station ids
for line in inputfile:
    if len(line.strip()) == 0:
        continue
    if line.startswith('#'):
        continue
    try: # comma separated
        origin, dest, count = line.split(',')
    except ValueError: # try tab separated, but otherwise give up
        try:
            origin, dest, count = line.split()
        except ValueError:
            warnings.warn("Unable to make sense of: %s" % line)
            continue

    origin, dest, count = int(origin), int(dest), int(count)
    trips[(origin, dest)] = count
    stations.add(origin)
    stations.add(dest)

total_trips = sum(trips.itervalues())

# Give some summary statistics
outfile.write("# Passenger file autogenerated with: \n# python %s\n" \
          % (' '.join(sys.argv)))

outfile.write("\n# Total passengers: %d arriving in %s seconds\n" \
          % (total_trips, options.endtime))

# Print a OD matrix
outfile.write("\n# Origin-Destination Matrix (Origin on Y, Dest on X)\n")
outfile.write("# -------------------------\n")

# header
station_list = list(stations)
station_list.sort()
outfile.write("#      " + "%5d"*len(stations) % tuple(station_list) + "\n")

# print matrix contents
for row in station_list:
    outfile.write("# " + "%5d" % row + "%5d"*len(stations) % tuple([trips[(row, col)] for col in station_list]) + "\n")

# generate the passengers
trip_list = []
for (orig, dest), cnt in trips.iteritems():
    for x in range(cnt):
        time = options.endtime * random.betavariate(*options.beta)
        trip_list.append( (time, orig, dest) )

trip_list.sort() # sorts by time

# header
outfile.write("\n\n# pID\torigin\tdest\ttime\tload\tunload\n")

# the trips
for pid, (time, orig, dest) in enumerate(trip_list):
    outfile.write("%s\t%5d\t%5d\t%10.3f\n" % (pid, orig, dest, time))